import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import pandas as pd
import tushare as ts
from typing import List
from utils.confutils import ConfUtil
import datetime
import time

def download_fin_report(fin_type:str,
                        period:str,
                        engine:ts.pro.client.DataApi,
                        report_type='1',
                        ):
    assert fin_type in ['cashflow','balancesheet','income'],f"fin_type 必须是 ['cashflow','balancesheet','income'] 中的一个，当前为{fin_type}"
    assert period[-4:] in ['1231','0930','0630','0331'],f"period最后四位数字必须是['1231','0930','0630','0331']之中的一个，当前为{period[-4:]}"
    if fin_type == 'income':
        
        download_func = engine.income_vip
    elif fin_type == "balancesheet":
        download_func = engine.balancesheet_vip
    else:
        download_func = engine.cashflow_vip
    
    df = download_func(period=period,report_type=report_type)
    return df

def download_all_fin_report(fin_type:str,
                            periods:List[str],
                            engine:ts.pro.client.DataApi,
                            report_type:str="1",
                            ):
    
    assert fin_type in ['cashflow','balancesheet','income'],f"fin_type 必须是 ['cashflow','balancesheet','income'] 中的一个，当前为{fin_type}"
    df_full = []
    print(f'开始下载{fin_type}')
    for period in periods:
        print(f'开始下载{period}全市场{fin_type}')
        while 1==1:
            try:
                df =  download_fin_report(fin_type,period,engine=engine,report_type=report_type)
                break
            except:
                time.sleep(5)
                continue
        
        df_full.append(df)
        time.sleep(1)
    print(f'{fin_type}下载完毕')
    res_df = pd.concat(df_full,axis=0)
    print(f'数据信息：\n{res_df.info()}')
    print(res_df.head(5))
    print('开始储存文件')
    tofile = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + f"\\Data\\RawData\\stock_{fin_type}.feather"
    res_df.to_feather(tofile)
    print('文件储存完毕')



    
if __name__ == "__main__":
    # 提取token
    token = ConfUtil().get_tokens('tushare')
    pro = ts.pro_api(token)
    # 下载股票财务数据,时间范围从20091231到当前日期
    today_str = datetime.datetime.now().strftime('%Y%m%d')
    periods = [i.strftime('%Y%m%d') for i in pd.date_range(start='20091231',end=today_str,freq='Q')]
    # 开始下载数据
    for fin in ['balancesheet','cashflow','income']:
        download_all_fin_report(fin,periods,engine=pro)
